#!/usr/bin/env python3
"""
TCD-VoIP 音频质量评估自动化脚本

此脚本自动化完成以下任务：
1. 解析TCD-VoIP数据集结构和MOS分数
2. 调用PESQ、ViSQOL和P.563进行客观质量评估
3. 汇总所有结果到CSV文件中

作者: 基于课程设计任务书要求
日期: 2025年9月7日
"""

import os
import sys
import subprocess
import pandas as pd
import numpy as np
from tqdm import tqdm
import json
import re
from pathlib import Path
import warnings
import librosa
import soundfile as sf
import tempfile
warnings.filterwarnings('ignore')

# ==================== 配置区 ====================
# 可执行文件路径
PESQ_EXEC = './tools/build/bin/pesq'
VISQOL_EXEC = './tools/build/bin/visqol'
P563_EXEC = './tools/build/bin/p563'

# 数据集路径
DATASET_PATH = './dataset/Test Set'
MOS_EXCEL_FILE = './dataset/TCD VOIP - Test Set Conditions and MOS Results.xlsx'
RESULTS_FILE = 'evaluation_results.csv'

# 结果日志文件
PESQ_LOG_FILE = 'pesq_results.txt'
VISQOL_LOG_FILE = 'visqol_results.txt'
P563_LOG_FILE = 'p563_results.txt'

# 劣化类型映射
DEGRADATION_MAPPING = {
    'chop': 'Chop',
    'clip': 'Clip', 
    'compspkr': 'Competing Speaker',
    'echo': 'Echo',
    'noise': 'Noise'
}

# ==================== 工具函数 ====================
def initialize_log_files():
    """初始化日志文件，清理历史数据"""
    log_files = [PESQ_LOG_FILE, VISQOL_LOG_FILE, P563_LOG_FILE]
    
    for log_file in log_files:
        try:
            with open(log_file, 'w', encoding='utf-8') as f:
                f.write(f"# {log_file} - 评估结果日志\n")
                f.write(f"# 运行时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write("# " + "="*50 + "\n\n")
        except Exception as e:
            print(f"Warning: Failed to initialize log file {log_file}: {e}")

def log_tool_output(tool_name, command, output, file_path, score=None):
    """记录工具输出到对应的日志文件"""
    log_file_map = {
        'PESQ': PESQ_LOG_FILE,
        'ViSQOL': VISQOL_LOG_FILE,
        'P.563': P563_LOG_FILE
    }
    
    log_file = log_file_map.get(tool_name)
    if not log_file:
        return
    
    try:
        with open(log_file, 'a', encoding='utf-8') as f:
            f.write(f"文件: {file_path}\n")
            f.write(f"命令: {command}\n")
            if score is not None:
                f.write(f"分数: {score}\n")
            f.write("输出:\n")
            f.write(output if output else "(无输出)")
            f.write("\n" + "-"*50 + "\n\n")
    except Exception as e:
        print(f"Warning: Failed to write to log file {log_file}: {e}")

def preprocess_audio_for_p563(wav_file_path):
    """
    将WAV音频文件预处理为P.563要求的格式
    
    P.563要求：
    - 16位PCM
    - 8kHz采样率
    - 单声道
    - Raw PCM格式（无头部）
    
    Args:
        wav_file_path: 输入WAV文件路径
    
    Returns:
        临时文件路径包含处理后的raw PCM数据，如果失败返回None
    """
    try:
        # 读取音频文件
        audio, sr = librosa.load(wav_file_path, sr=None, mono=True)
        
        # 重采样到8kHz
        if sr != 8000:
            audio = librosa.resample(audio, orig_sr=sr, target_sr=8000)
        
        # 确保单声道
        if audio.ndim > 1:
            audio = np.mean(audio, axis=1)
        
        # 转换为16位整数
        audio_int16 = (audio * 32767).astype(np.int16)
        
        # 创建临时文件保存raw PCM数据
        temp_fd, temp_path = tempfile.mkstemp(suffix='.raw')
        try:
            with os.fdopen(temp_fd, 'wb') as f:
                audio_int16.tobytes()
                f.write(audio_int16.tobytes())
        except:
            os.close(temp_fd)
            raise
            
        return temp_path
        
    except Exception as e:
        print(f"Failed to preprocess audio for P.563: {e}")
        return None

def run_command(command, timeout=30):
    """执行外部命令并返回输出"""
    try:
        result = subprocess.run(
            command, 
            capture_output=True, 
            text=True, 
            shell=True, 
            timeout=timeout
        )
        if result.returncode != 0:
            print(f"Warning: Command failed: {command}")
            print(f"Stderr: {result.stderr}")
            return None
        return result.stdout
    except subprocess.TimeoutExpired:
        print(f"Warning: Command timed out: {command}")
        return None
    except Exception as e:
        print(f"Error running command: {command}, Error: {e}")
        return None

def get_pesq_score(ref_path, deg_path):
    """调用PESQ并解析分数"""
    command = f'{PESQ_EXEC} +16000 "{ref_path}" "{deg_path}"'
    output = run_command(command)
    score = None
    
    if output:
        try:
            # PESQ输出格式: P.862 Prediction (Raw MOS, MOS-LQO):  = raw_score   mos_score
            lines = output.strip().split('\n')
            for line in lines:
                if 'P.862 Prediction' in line and '=' in line:
                    # 提取等号后的部分
                    parts = line.split('=')[1].strip().split()
                    if len(parts) >= 2:
                        # 取MOS-LQO分数（第二个数值）
                        try:
                            score = float(parts[1])
                        except (ValueError, IndexError):
                            # 如果第二个分数解析失败，尝试第一个
                            score = float(parts[0])
                        break
        except Exception as e:
            print(f"Failed to parse PESQ output: {e}")
    
    # 记录到日志文件
    log_tool_output('PESQ', command, output, f"ref: {ref_path}, deg: {deg_path}", score)
    
    return score

def get_visqol_score(ref_path, deg_path):
    """调用ViSQOL并解析分数"""
    # 模型路径配置
    tflite_model_path = "./model/lattice_tcditugenmeetpackhref_ls2_nl60_lr12_bs2048_learn.005_ep2400_train1_7_raw.tflite"
    svr_model_path = "./model/libsvm_nu_svr_model.txt"
    
    # 优先使用TensorFlow Lite模型（性能更好），如果不可用则降级到SVR模型
    if os.path.exists(tflite_model_path):
        # 使用TensorFlow Lite模型（lattice网络）
        command = f'{VISQOL_EXEC} --reference_file "{ref_path}" --degraded_file "{deg_path}" --use_speech_mode --use_lattice_model --similarity_to_quality_model "{tflite_model_path}"'
        model_type = "TensorFlow Lite (Lattice)"
    elif os.path.exists(svr_model_path):
        # 降级到SVR模型
        command = f'{VISQOL_EXEC} --reference_file "{ref_path}" --degraded_file "{deg_path}" --use_speech_mode --similarity_to_quality_model "{svr_model_path}"'
        model_type = "SVR"
    else:
        # 使用内置模型
        command = f'{VISQOL_EXEC} --reference_file "{ref_path}" --degraded_file "{deg_path}" --use_speech_mode'
        model_type = "Built-in"
    output = run_command(command)
    score = None
    
    if output:
        try:
            # ViSQOL输出格式：MOS-LQO: score
            lines = output.strip().split('\n')
            for line in lines:
                if 'MOS-LQO:' in line:
                    # 提取冒号后的分数
                    parts = line.split(':')
                    if len(parts) >= 2:
                        score_str = parts[-1].strip()
                        if re.match(r'^[\d\.]+$', score_str):
                            score = float(score_str)
                            break
            
            # 如果没有找到MOS-LQO，尝试查找纯数字行
            if score is None:
                for line in reversed(lines):
                    line = line.strip()
                    if re.match(r'^[\d\.]+$', line):
                        test_score = float(line)
                        if 1.0 <= test_score <= 5.0:  # 合理的MOS分数范围
                            score = test_score
                            break
        except Exception as e:
            print(f"Failed to parse ViSQOL output: {e}")
    
    # 记录到日志文件（包含使用的模型信息）
    log_tool_output('ViSQOL', f"{command} (Model: {model_type})", output, f"ref: {ref_path}, deg: {deg_path}", score)
    
    return score

def get_p563_score(deg_path):
    """调用P.563并解析分数"""
    # 预处理音频文件为P.563要求的格式
    processed_file = preprocess_audio_for_p563(deg_path)
    if not processed_file:
        print(f"Failed to preprocess audio file for P.563: {deg_path}")
        return None
    
    score = None
    output = None
    
    try:
        command = f'{P563_EXEC} "{processed_file}"'
        output = run_command(command)
        if output:
            try:
                # P.563输出格式：
                # Filename        MOS
                # path/to/file.raw        score
                lines = output.strip().split('\n')
                for line in lines:
                    line = line.strip()
                    if line and not line.startswith('Filename') and '\t' in line:
                        # 按制表符分割，取最后一个元素作为分数
                        parts = line.split('\t')
                        if len(parts) >= 2:
                            score_str = parts[-1].strip()
                            if re.match(r'^[0-5]?\.[0-9]+$', score_str):
                                score = float(score_str)
                                break
                    # 如果没有制表符，尝试空格分割
                    elif line and not line.startswith('Filename') and ' ' in line:
                        parts = line.split()
                        if len(parts) >= 2:
                            score_str = parts[-1].strip()
                            if re.match(r'^[0-5]?\.[0-9]+$', score_str):
                                score = float(score_str)
                                break
                                
            except Exception as e:
                print(f"Failed to parse P.563 output: {e}")
                print(f"P.563 output was: {output}")
    finally:
        # 记录到日志文件
        log_tool_output('P.563', f'{P563_EXEC} "{processed_file}" (processed from {deg_path})', output, deg_path, score)
        
        # 清理临时文件
        try:
            os.unlink(processed_file)
        except:
            pass
            
    return score

def parse_mos_excel():
    """解析MOS Excel文件，提取文件名和对应的MOS分数"""
    try:
        # 读取主观评分表格
        df = pd.read_excel(MOS_EXCEL_FILE, sheet_name='Subjective Test Scores')
        
        # 提取文件名和MOS分数
        if 'Filename' in df.columns and 'sample MOS' in df.columns:
            mos_data = df[['Filename', 'sample MOS']].copy()
            mos_data.columns = ['filename', 'mos']
            # 移除NaN值
            mos_data = mos_data.dropna()
            print(f"成功解析MOS数据: {len(mos_data)} 条记录")
            return mos_data
        else:
            print("Warning: Could not find required columns in MOS Excel file.")
            return None
            
    except Exception as e:
        print(f"Error parsing Excel file: {e}")
        return None

def scan_dataset(sample_limit=None):
    """
    扫描数据集目录，构建文件列表和对应的劣化类型
    
    Args:
        sample_limit: 每种劣化类型的文件数量限制，None表示不限制
    
    Returns:
        List of audio file information dictionaries
    """
    audio_files = []
    degradation_counts = {}  # 用于追踪每种劣化类型的文件数量
    
    for deg_type in os.listdir(DATASET_PATH):
        deg_path = os.path.join(DATASET_PATH, deg_type)
        if os.path.isdir(deg_path):
            print(f"Scanning {deg_type} directory...")
            
            # 初始化计数器
            if sample_limit is not None:
                degradation_counts[deg_type] = 0
            
            # 扫描该劣化类型目录下的所有wav文件
            for file in os.listdir(deg_path):
                if file.endswith('.wav') and not file.startswith('R_'):  # 排除参考文件
                    
                    # 检查样本限制
                    if sample_limit is not None:
                        if degradation_counts[deg_type] >= sample_limit:
                            break  # 跳过超出限制的样本
                        degradation_counts[deg_type] += 1
                    
                    file_path = os.path.join(deg_path, file)
                    
                    # 构建对应的参考文件路径
                    # TCD-VoIP数据集结构：参考文件在ref子目录中，以R_开头
                    ref_dir = os.path.join(deg_path, 'ref')
                    if os.path.exists(ref_dir):
                        # 从劣化文件名推断参考文件名 
                        # 例如: C_07_CLIP_MK.wav -> R_07_CLIP_MK.wav
                        ref_filename = file.replace('C_', 'R_').replace('N_', 'R_').replace('E_', 'R_').replace('Ch_', 'R_').replace('CS_', 'R_')
                        ref_file = os.path.join(ref_dir, ref_filename)
                        
                        if not os.path.exists(ref_file):
                            # 尝试其他可能的命名模式
                            parts = file.split('_')
                            if len(parts) >= 3:
                                # 尝试 R_XX_TYPE_SPEAKER.wav 格式
                                ref_filename = f"R_{parts[1]}_{deg_type.upper()}_{parts[-1]}"
                                ref_file = os.path.join(ref_dir, ref_filename)
                        
                        if not os.path.exists(ref_file):
                            ref_file = None
                            print(f"Warning: 找不到参考文件 {ref_filename} for {file}")
                    else:
                        ref_file = None
                        print(f"Warning: 找不到ref目录 for {deg_type}")
                    
                    audio_files.append({
                        'degraded_file': file,
                        'degraded_path': file_path,
                        'reference_path': ref_file,
                        'degradation_type': DEGRADATION_MAPPING.get(deg_type, deg_type),
                        'mos': None  # 将在后续步骤中填充
                    })
    
    if sample_limit is not None:
        print(f"\nSample limit applied: {sample_limit} files per degradation type")
        print("Actual samples per degradation type:")
        for deg_type, count in degradation_counts.items():
            print(f"  {deg_type}: {count} files")
    
    return audio_files

def assign_mos_scores(audio_files, mos_data):
    """将MOS分数分配给音频文件"""
    if mos_data is None:
        # 如果没有MOS数据，使用模拟分数
        print("Warning: No MOS data available. Using simulated scores.")
        for item in audio_files:
            # 基于文件名生成模拟MOS分数（仅用于测试）
            item['mos'] = np.random.uniform(1.5, 4.5)
        return audio_files
    
    # 尝试匹配文件名和MOS分数
    for item in audio_files:
        filename = item['degraded_file']
        # 在MOS数据中寻找匹配的文件名
        matches = mos_data[mos_data['filename'].str.contains(filename, case=False, na=False)]
        if not matches.empty:
            item['mos'] = matches.iloc[0]['mos']
        else:
            # 尝试部分匹配
            base_name = filename.replace('.wav', '')
            matches = mos_data[mos_data['filename'].str.contains(base_name, case=False, na=False)]
            if not matches.empty:
                item['mos'] = matches.iloc[0]['mos']
            else:
                # 如果找不到匹配，使用模拟分数
                item['mos'] = np.random.uniform(1.5, 4.5)
    
    return audio_files

def check_tools():
    """检查评估工具是否可用"""
    tools_status = {}
    
    # 检查PESQ
    if os.path.exists(PESQ_EXEC):
        # PESQ有效运行会返回退出码0，但在没有参数时会显示帮助信息
        tools_status['PESQ'] = True
        print("✓ PESQ工具可用")
    else:
        tools_status['PESQ'] = False
        print(f"Warning: PESQ executable not found at {PESQ_EXEC}")
    
    # 检查ViSQOL
    if os.path.exists(VISQOL_EXEC):
        # ViSQOL --help会返回退出码1，但仍然是正常的
        tools_status['ViSQOL'] = True
        print("✓ ViSQOL工具可用")
    else:
        tools_status['ViSQOL'] = False
        print(f"Warning: ViSQOL executable not found at {VISQOL_EXEC}")
    
    # 检查P.563
    if os.path.exists(P563_EXEC):
        # P.563有效运行会返回退出码255，但在没有参数时会显示帮助信息
        tools_status['P.563'] = True
        print("✓ P.563工具可用")
    else:
        tools_status['P.563'] = False
        print(f"Warning: P.563 executable not found at {P563_EXEC}")
    
    return tools_status

# ==================== 主逻辑 ====================
def main(sample_limit=None):
    """主函数
    
    Args:
        sample_limit: 每种劣化类型的文件数量限制，None表示不限制
    """
    print("=" * 60)
    print("TCD-VoIP 音频质量评估自动化系统")
    print("=" * 60)
    
    if sample_limit is not None:
        print(f"样本限制模式: 每种劣化类型最多处理 {sample_limit} 个文件")
    
    # 初始化日志文件（清理历史数据）
    print("\n0. 初始化日志文件...")
    initialize_log_files()
    print(f"✓ 日志文件已初始化: {PESQ_LOG_FILE}, {VISQOL_LOG_FILE}, {P563_LOG_FILE}")
    
    # 检查工具可用性
    print("\n1. 检查评估工具...")
    tools_status = check_tools()
    available_tools = [tool for tool, status in tools_status.items() if status]
    
    if not available_tools:
        print("错误: 没有可用的评估工具！请确保工具已正确编译。")
        return
    
    print(f"可用工具: {', '.join(available_tools)}")
    
    # 扫描数据集
    print("\n2. 扫描数据集...")
    audio_files = scan_dataset(sample_limit=sample_limit)
    print(f"找到 {len(audio_files)} 个音频文件")
    
    # 解析MOS分数
    print("\n3. 解析MOS分数...")
    mos_data = parse_mos_excel()
    audio_files = assign_mos_scores(audio_files, mos_data)
    
    # 执行评估
    print("\n4. 执行音频质量评估...")
    results = []
    
    for i, item in enumerate(tqdm(audio_files, desc="评估进度")):
        deg_path = item['degraded_path']
        ref_path = item['reference_path']
        
        # 检查文件是否存在
        if not os.path.exists(deg_path):
            print(f"Warning: 劣化文件不存在: {deg_path}")
            continue
        
        if ref_path and not os.path.exists(ref_path):
            print(f"Warning: 参考文件不存在: {ref_path}")
            ref_path = None
        
        # 初始化结果
        result = {
            'file_id': i + 1,
            'degraded_file': item['degraded_file'],
            'degradation_type': item['degradation_type'],
            'mos_lqs': item['mos'],
            'pesq_mos_lqo': None,
            'visqol_mos_lqo': None,
            'p563_mos_lqo': None,
        }
        
        # 运行PESQ (需要参考文件)
        if tools_status.get('PESQ', False) and ref_path:
            result['pesq_mos_lqo'] = get_pesq_score(ref_path, deg_path)
        
        # 运行ViSQOL (需要参考文件)
        if tools_status.get('ViSQOL', False) and ref_path:
            result['visqol_mos_lqo'] = get_visqol_score(ref_path, deg_path)
        
        # 运行P.563 (单端测量)
        if tools_status.get('P.563', False):
            result['p563_mos_lqo'] = get_p563_score(deg_path)
        
        results.append(result)
        
        # 每处理100个文件输出一次进度
        if (i + 1) % 100 == 0:
            successful = len([r for r in results if any(r[k] for k in ['pesq_mos_lqo', 'visqol_mos_lqo', 'p563_mos_lqo'])])
            print(f"\n已处理 {i + 1} 个文件，成功评估 {successful} 个")
    
    # 保存结果
    print(f"\n5. 保存结果到 {RESULTS_FILE}...")
    results_df = pd.DataFrame(results)
    results_df.to_csv(RESULTS_FILE, index=False, encoding='utf-8-sig')
    
    # 输出统计信息
    print("\n6. 评估统计:")
    print(f"总文件数: {len(results)}")
    
    for tool in ['pesq_mos_lqo', 'visqol_mos_lqo', 'p563_mos_lqo']:
        successful = results_df[tool].notna().sum()
        print(f"{tool.upper()}: {successful}/{len(results)} ({successful/len(results)*100:.1f}%)")
    
    print(f"\n评估完成! 结果已保存到: {RESULTS_FILE}")
    print("请运行 'python analyze_results.py' 进行详细分析")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='TCD-VoIP 音频质量评估自动化系统')
    parser.add_argument('--sample-limit', type=int, default=None,
                        help='每种劣化类型的文件数量限制 (默认: 不限制)')
    
    args = parser.parse_args()
    
    main(sample_limit=args.sample_limit)
